//
// Created by zy on 22-6-7.
//
#include <cassert>
#include <chrono>
#include <iostream>
#include "leveldb/db.h"
#include "leveldb/comparator.h"
#include "util.h"
#include "stats.h"
#include "learned_index.h"
#include <cstring>
#include "cxxopts.hpp"
#include <unistd.h>
#include <fstream>
#include <string>
#include <sstream>
#include "../db/version_set.h"
#include <cmath>
#include <random>

using namespace std::chrono;
using namespace leveldb;
using namespace adgMod;
using std::string;
using std::cerr;
using std::cout;
using std::endl;
using std::to_string;
using std::vector;
using std::map;
using std::ifstream;
using std::string;

int num_pairs_base = 1000;
int mix_base = 20;
#include <set>


class NumericalComparator : public Comparator {
 public:
  NumericalComparator() = default;
  virtual const char* Name() const {return "adgMod:NumericalComparator";}
  virtual int Compare(const Slice& a, const Slice& b) const {
    uint64_t ia = adgMod::ExtractInteger(a.data(), a.size());
    uint64_t ib = adgMod::ExtractInteger(b.data(), b.size());
    if (ia < ib) return -1;
    else if (ia == ib) return 0;
    else return 1;
  }
  virtual void FindShortestSeparator(std::string* start, const Slice& limit) const { return; };
  virtual void FindShortSuccessor(std::string* key) const { return; };
};


void PutAndPrefetch(int lower, int higher, vector<string>& keys) {
  adgMod::Stats* instance = adgMod::Stats::GetInstance();

  Status status;

  instance->StartTimer(9);
  for (int i = lower; i < higher; ++i) {
    string value = generate_value(0);
    status = db->Put(write_options, keys[i], value);
    assert(status.ok() && "File Put Error");
  }
  instance->PauseTimer(9, true);

  //cout << "Put Complete" << endl;


  instance->StartTimer(10);
  for (int i = lower; i < higher; ++i) {
    string value;
    status = db->Get(read_options, keys[i], &value);
    //cout << "Get " << keys[i] << " Done" << endl;
    assert(status.ok() && "File Get Error");
  }
  instance->PauseTimer(10, true);

  //cout << "Prefetch Complete" << endl;
};


int main(int argc, char *argv[]) {
  int num_operations, num_iteration, num_mix;
  float test_num_segments_base;
  float num_pair_step;
  string db_location, profiler_out, input_filename, distribution_filename,
      ycsb_filename;
  bool print_single_timing, print_file_info, evict, unlimit_fd,
      use_distribution = false, pause, use_ycsb = false;
  bool change_level_load, change_file_load, change_level_learning,
      change_file_learning;
  int insert_bound, length_range;
  string db_location_copy;

  cxxopts::Options commandline_options("leveldb read test",
                                       "Testing leveldb read performance.");
  commandline_options.add_options()(
      "n,get_number", "the number of gets (to be multiplied by 1024)",
      cxxopts::value<int>(num_operations)->default_value("1000"))(
      "s,step", "the step of the loop of the size of db",
      cxxopts::value<float>(num_pair_step)->default_value("1"))(
      "i,iteration", "the number of iterations of a same size",
      cxxopts::value<int>(num_iteration)->default_value("1"))(
      "m,modification", "if set, run our modified version",
      cxxopts::value<int>(adgMod::MOD)->default_value("0"))(
      "h,help", "print help message",
      cxxopts::value<bool>()->default_value("false"))(
      "d,directory", "the directory of db",
      cxxopts::value<string>(db_location)->default_value("/mnt/ssd/testdb"))(
      "k,key_size", "the size of key",
      cxxopts::value<int>(adgMod::key_size)->default_value("8"))(
      "v,value_size", "the size of value",
      cxxopts::value<int>(adgMod::value_size)->default_value("8"))(
      "single_timing", "print the time of every single get",
      cxxopts::value<bool>(print_single_timing)->default_value("false"))(
      "file_info", "print the file structure info",
      cxxopts::value<bool>(print_file_info)->default_value("false"))(
      "test_num_segments", "test: number of segments per level",
      cxxopts::value<float>(test_num_segments_base)->default_value("1"))(
      "string_mode", "test: use string or int in model",
      cxxopts::value<bool>(adgMod::string_mode)->default_value("false"))(
      "e,model_error", "error in modesl",
      cxxopts::value<uint32_t>(adgMod::model_error)->default_value("8"))(
      "f,input_file", "the filename of input file",
      cxxopts::value<string>(input_filename)->default_value(""))(
      "multiple", "test: use larger keys",
      cxxopts::value<uint64_t>(adgMod::key_multiple)->default_value("1"))(
      "w,write", "writedb",
      cxxopts::value<bool>(fresh_write)->default_value("false"))(
      "c,uncache", "evict cache",
      cxxopts::value<bool>(evict)->default_value("false"))(
      "u,unlimit_fd", "unlimit fd",
      cxxopts::value<bool>(unlimit_fd)->default_value("false"))("x,dummy",
                                                                "dummy option")(
      "filter", "use filter",
      cxxopts::value<bool>(adgMod::use_filter)->default_value("false"))(
      "mix", "mix read and write",
      cxxopts::value<int>(num_mix)->default_value("0"))(
      "distribution", "operation distribution",
      cxxopts::value<string>(distribution_filename)->default_value(""))(
      "change_level_load", "load level model",
      cxxopts::value<bool>(change_level_load)->default_value("false"))(
      "change_file_load", "enable level learning",
      cxxopts::value<bool>(change_file_load)->default_value("false"))(
      "change_level_learning", "load file model",
      cxxopts::value<bool>(change_level_learning)->default_value("false"))(
      "change_file_learning", "enable file learning",
      cxxopts::value<bool>(change_file_learning)->default_value("false"))(
      "p,pause", "pause between operation",
      cxxopts::value<bool>(pause)->default_value("false"))(
      "policy", "learn policy",
      cxxopts::value<int>(adgMod::policy)->default_value("0"))(
      "YCSB", "use YCSB trace",
      cxxopts::value<string>(ycsb_filename)->default_value(""))(
      "insert", "insert new value",
      cxxopts::value<int>(insert_bound)->default_value("0"))(
      "range", "use range query and specify length",
      cxxopts::value<int>(length_range)->default_value("0"));
  auto result = commandline_options.parse(argc, argv);
  if (result.count("help")) {
    printf("%s", commandline_options.help().c_str());
    exit(0);
  }

  std::default_random_engine e1(0), e2(255);
  srand(0);
  num_operations *= num_pairs_base;  // 1000 * 1000,什么意思呢？

  adgMod::fd_limit = unlimit_fd ? 1024 * 1024 : 1024;  // TODO: 需要测试一下是否真的影响性能
  adgMod::restart_read = true;
  adgMod::level_learning_enabled ^= change_level_learning;
  adgMod::file_learning_enabled ^= change_file_learning;
  adgMod::load_level_model ^= change_level_load;
  adgMod::load_file_model ^= change_file_load;

  // adgMod::file_learning_enabled = false;

  vector<uint64_t> distribution;  // uint64_t是unsigned long int的typedef别名
  vector<int> ycsb_is_write;
  // keys.reserve(100000000000 / adgMod::value_size);

//  std::cerr << "读取keys from " << input_filename << std::endl;
  ifstream input(input_filename);
  assert(input.good());  // 加个判定，提前结束错误的情况
  std::string dev_id;
  char* op_flag;
  std::string key;
  vector<string> keys_all;
  vector<char> ops_all;

  string line;
  while (getline(input, line)) {
    vector<string> data_line;
    string number;  // 一行里的每个用逗号分隔的部分
    std::istringstream readstr(line);
    // 将一行数据按逗号分割
    for (int j=0; j<5; j++) {
      getline(readstr, number, ',');
      data_line.push_back(number);
    }
    op_flag = (char*) (data_line[1]).data();
    key = data_line[2];
    string the_key = generate_key(key);
    keys_all.push_back(std::move(the_key));
    ops_all.push_back(std::move(*op_flag));
  }

  int N = int(keys_all.size() * 3.0/5.0);

  std::vector<std::string> keys_former(keys_all.begin(), keys_all.begin() + N);
  std::vector<std::string> keys_latter(keys_all.begin() + N, keys_all.end());
  std::vector<char> ops_former(ops_all.begin(), ops_all.begin() + N);
  std::vector<char> ops_latter(ops_all.begin() + N, ops_all.end());
  assert(keys_former.size() == ops_former.size());
  keys_all.clear();
  ops_all.clear();

  int ops_former_all_num = ops_former.size();
  int ops_former_read_num = get_read_nums(ops_former);
  int ops_former_write_num = ops_former_all_num - ops_former_read_num;

  int ops_latter_all_num = ops_latter.size();
  int ops_latter_read_num = get_read_nums(ops_latter);
  int ops_latter_write_num = ops_latter_all_num - ops_latter_read_num;

  std::cout << "formerall-" << std::to_string(ops_former_all_num) << "-formerread-" << std::to_string(ops_former_read_num) << "-formerwrite-" << std::to_string(ops_former_write_num) << "-" << std::endl;
  std::cout << "latterall-" << std::to_string(ops_latter_all_num) << "-latterread-" << std::to_string(ops_latter_read_num) << "-latterwrite-" << std::to_string(ops_latter_write_num) << "-" << std::endl;

  // 在load_ar的时候copy_out是false
  adgMod::Stats* instance = adgMod::Stats::GetInstance();
  vector<vector<size_t>> times(20);
  string values(1024 * 1024, '0');


  for (size_t iteration = 0; iteration < num_iteration; ++iteration) {
    std::cerr << "allPut, iteration: " << iteration <<std::endl;
    // 随机有关的
    std::uniform_int_distribution<uint64_t> uniform_dist_value(0, (uint64_t)values.size() - adgMod::value_size - 1);

    // 定义数据库指针，数据库操作选项。在循环里定义可能是每次要wait，delete一下，好像更符合纯测查找？
    DB* db;
    Options options;
    ReadOptions& read_options = adgMod::read_options;
    WriteOptions& write_options = adgMod::write_options;
    Status status;
    options.create_if_missing = true;
    write_options.sync = false;

    // instance必须每个iteration都重置
    instance->ResetAll();

    if (iteration == 0) {
      fresh_write = true;  // 以防万一
      // 写入开始
      string command = "rm -rf " + db_location;
      system(command.c_str());
      system("sudo fstrim -a -v");
      system("sync; echo 3 | sudo tee /proc/sys/vm/drop_caches");
      cout << "delete and trim complete" << endl;

      status = DB::Open(options, db_location, &db);
      assert(status.ok() && "Open Error");

      instance->StartTimer(9);
      // 加载数据库不要chunk了，因为根本不打乱
      for (int i=0; i<keys_former.size(); i++) {
        if (ops_former[i] == 'W') {
          status = db->Put(write_options, keys_former[i], {values.data() + uniform_dist_value(e2), (uint64_t)adgMod::value_size});
          assert(status.ok() && "File Put Error");
        } else if (ops_former[i] == 'R') {
          std::string value_load;
          db->Get(read_options, keys_former[i], &value_load);
          // 这里不是测试部分，只是为了模拟加载数据库，所以不断言
        } else throw "wrong operation flag";
      }
      adgMod::db->vlog->Sync();
      instance->PauseTimer(9, true);
      std::cout << "Put Complete" << std::endl;

      if (print_file_info && iteration==0)  db->PrintFileInfo();
      adgMod::db->WaitForBackground();
      delete db;

      // 第二阶段，开始建立学习索引。注意建立学习索引是跟加载数据库一一对应的
      time_point<system_clock> start_training = system_clock::now();
      status = DB::Open(options, db_location, &db);
      assert(status.ok() && "second open error");
      adgMod::db->WaitForBackground();

      if (adgMod::MOD == 6 || adgMod::MOD == 7) {
        Version* current = adgMod::db->versions_->current();
        for (int i = 1; i < config::kNumLevels; ++i) {
          LearnedIndexData::Learn(new VersionAndSelf{current, adgMod::db->version_count, current->learned_index_data_[i].get(), i});
        }
        current->FileLearn();
      }

      std::cout << "Shutting down" << std::endl;
      adgMod::db->WaitForBackground();
      delete db;
      time_point<system_clock> end_training = system_clock::now();
      std::chrono::duration<double> elapsed_training  = end_training - start_training;
      std::cout << "iteration: " << iteration << ", training time: " << elapsed_training.count() << " sec" << std::endl;

      std::string data_name_4print = get_dataname_from_filepath(input_filename);
      std::cerr << data_name_4print << "load and learn finished" << std::endl;

      // 开始进入下一阶段，非fresh_write
      fresh_write = false;
//      break;
    }

    // 加载完之后，缓存淘汰，可以用来对比
    if (evict) {
      std::cerr << "drop cache" << endl;
      system("sync; echo 3 | sudo tee /proc/sys/vm/drop_caches");
    }

    // 第三阶段了，开始测试。作为第一轮后半部分和2-5轮的全部部分。
    // data_name各个period都一样，每个iteration保存sstable的信息
    std::string data_name = get_dataname_from_filepath(input_filename);
    // 另一种指定方式，N_Period为后40%的1/5。
    int N_Period = int(0.02 * keys_latter.size());
    readFileList((char*) db_location.data(), data_name, adgMod::MOD, "trace_allPut", iteration, adgMod::key_size);

    std::vector<std::string> keys_latter_inPeriod;
    std::vector<char> ops_latter_inPeriod;
    if ((iteration+1) * N_Period <= keys_latter.size()) {
      keys_latter_inPeriod.assign(keys_latter.begin() + iteration * N_Period, keys_latter.begin() + (iteration+1) * N_Period);
      ops_latter_inPeriod.assign(ops_latter.begin() + iteration * N_Period, ops_latter.begin() + (iteration+1) * N_Period);
      int op_all_num = ops_latter_inPeriod.size();
      int op_read_num = get_read_nums(ops_latter_inPeriod);
      int op_write_num = op_all_num - op_read_num;
      std::cout << "period-" << iteration << "-all-" << std::to_string(op_all_num) << "-read-" << std::to_string(op_read_num) << "-write-" << std::to_string(op_write_num) << "-" << std::endl;
    } else {
      std::cout << "period-" << iteration << "-all-" << "0" << "-" << std::endl;
      continue;
    }

    std::cout << "Starting up" << std::endl;
    status = DB::Open(options, db_location, &db);
    adgMod::db->WaitForBackground();
    assert(status.ok() && "Open Error");

    uint64_t last_read = 0, last_write = 0;
    int last_level = 0, last_file = 0, last_baseline = 0, last_succeeded = 0,
        last_false = 0, last_compaction = 0, last_learn = 0;
    std::vector<uint64_t> detailed_times;
    bool start_new_event = true;

    int num_operations = keys_latter_inPeriod.size();
    instance->StartTimer(13);

    for (int i = 0; i < num_operations; ++i) {
      if (start_new_event) {
        detailed_times.push_back(instance->GetTime());
        start_new_event = false;
      }

      uint64_t index = i;
      const string& key = keys_latter_inPeriod[index];

      // all operations are Put
      instance->StartTimer(10);
      status = db->Put(write_options, key, {values.data() + uniform_dist_value(e2), (uint64_t)64});
      instance->PauseTimer(10);
      assert(status.ok() && "Put Error");

      if (pause) {
        if ((i + 1) % (num_operations / 10000) == 0) ::usleep(800000);
      }

      if ((i + 1) % (num_operations / 100) == 0) detailed_times.push_back(instance->GetTime());
      if ((i + 1) % (num_operations / 10) == 0) {
        int level_read = levelled_counters[0].Sum();
        int file_read = levelled_counters[1].Sum();
        int baseline_read = levelled_counters[2].Sum();
        int succeeded_read = levelled_counters[3].NumSum();
        int false_read = levelled_counters[4].NumSum();

        compaction_counter_mutex.Lock();
        int num_compaction = events[0].size();
        compaction_counter_mutex.Unlock();
        learn_counter_mutex.Lock();
        int num_learn = events[1].size();
        learn_counter_mutex.Unlock();

        uint64_t read_time = instance->ReportTime(4);
        uint64_t write_time = instance->ReportTime(10);
        std::pair<uint64_t, uint64_t> time = {detailed_times.front(), detailed_times.back()};

        events[2].push_back(new WorkloadEvent(time, level_read - last_level, file_read - last_file,
            baseline_read - last_baseline, succeeded_read - last_succeeded,
            false_read - last_false, num_compaction - last_compaction,
            num_learn - last_learn, read_time - last_read,
            write_time - last_write, std::move(detailed_times)));

        last_level = level_read;
        last_file = file_read;
        last_baseline = baseline_read;
        last_succeeded = succeeded_read;
        last_false = false_read;
        last_compaction = num_compaction;
        last_learn = num_learn;
        last_read = read_time;
        last_write = write_time;
        detailed_times.clear();
        start_new_event = true;
        std::cout << (i + 1) / (num_operations / 10) << std::endl;
        Version* current = adgMod::db->versions_->current();
        printf("LevelSize %d %d %d %d %d %d\n", current->NumFiles(0),
               current->NumFiles(1), current->NumFiles(2),
               current->NumFiles(3), current->NumFiles(4),
               current->NumFiles(5));
      }
    }

    instance->PauseTimer(13, true);


    instance->ReportTime();
    for (int s = 0; s < times.size(); ++s) {
      times[s].push_back(instance->ReportTime(s));
    }
    adgMod::db->WaitForBackground();
    sleep(10);

    for (auto& event_array : events) {
      for (Event* e : event_array) e->Report();
    }

    for (Counter& c : levelled_counters) c.Report();

    file_data->Report();

    for (auto it : file_stats) {
      printf("FileStats %d %d %lu %lu %u %u %ld %d\n", it.first,
             it.second.level, it.second.start, it.second.end,
             it.second.num_lookup_pos, it.second.num_lookup_neg,
             it.second.size, it.first < file_data->watermark ? 0 : 1);
    }

    adgMod::learn_cb_model->Report();
    delete db;
  }  // 整个iteration结束


  std::string data_name_4print_out = get_dataname_from_filepath(input_filename);
  if (adgMod::MOD == 7) {
    std::cerr << data_name_4print_out << "-put-llsm done" << std::endl;
  } else {
    std::cerr << data_name_4print_out << "-put-bs done" << std::endl;
  }

  // 统计部分，打印到文件里面
  if (num_iteration > 1) {
    perf_exclude_1st_Iteration(times);
  }
}